import numpy
import torch
import torchvision
import torchvision.transforms.functional
import matplotlib.pyplot as plot
import torch.nn.functional as F
from torch import linalg as LA


class STLNet(torch.nn.Sequential):
    def __init__(self):
        super().__init__()
        self.backbone = torchvision.models.resnet50(pretrained=True)
        self.max_pool = torch.nn.Sequential(*list(self.backbone.children())[0:4])
        self.layer1 = self.backbone.layer1
        self.layer2 = self.backbone.layer2
        self.down_sample = torch.nn.Conv2d(in_channels=256, out_channels=256, kernel_size=(1, 1), stride=(2, 2))

    def forward(self, x):
        x = self.max_pool(x)
        x = self.layer1(x)
        x1 = self.layer2(x)
        x = self.down_sample(x)

        return torch.cat([x, x1], dim=1)


if __name__ == '__main__':
    img = torchvision.io.read_image('./resources/01.jpg').unsqueeze(0)
    img = torchvision.transforms.functional.resize(img, [224, 224]).float()

    model = STLNet()
    A = model(img)
    g = F.adaptive_avg_pool2d(A, (1, 1))
    A2 = LA.norm(A, dim=1, ord=2)
    g2 = LA.norm(g, dim=1, ord=2)
    S = g * A / (g2 * A2)
    print(S.shape)
